Adding KV-Cache Support to Flux2.Klein Image Editing - #465
Open
amepas wants to merge 2 commits into
Open
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for the Flux2Klein model, including new configuration files, NNX-based model implementations for the Transformer and VAE, and a dedicated pipeline. It also adds end-to-end parity and smoke tests. The review identified several critical issues: missing imports in flux2klein_pipeline.py and generate_flux2klein.py that will cause runtime errors, and the incorrect use of nnx.silu instead of jax.nn.silu across multiple model files.
…X (4B & 9B) with E2E parity and smoke tests - Add multi-image editing pipeline (FlaxFlux2KleinPipeline) supporting arbitrary reference image conditioning - Port VAE encoder and decoder to Flax NNX (AutoencoderKLFlux2NNX) - Implement efficient image preprocessing, tiling, aspect-ratio-aware padding, and packing utilities - Support fast sharded safetensors loading and TP sharding for Klein 4B & 9B - Use standard HF_HOME resolution without machine-specific hardcoded paths - Ensure clean compatibility across Transformers 4.x and 5.x via lazy module loading and dynamic FlaxPreTrainedModel lookup - Add end-to-end multi-image editing parity test and preprocessing unit test suite
amepas
force-pushed
the
onboarding-imageedit-kv-flux2klein
branch
2 times, most recently
from
August 26, 2026 16:43
6238524 to
f04bf0e
Compare
…age editing with E2E parity and smoke tests - Implement prefix extraction phase (step 0) and cached denoising scan loop in FlaxFlux2KleinPipeline - Support KV cache slicing and concatenation across Flax NNX Double and Single transformer blocks - Safely bound Splash/Flash attention block sizes for asymmetric cross-attention sequences with opt-in flag - Add CLI and config support for use_kv with dynamic FLUX.2-klein-9b-kv repository resolution - Add cross-framework E2E parity test achieving 0.8881 SSIM / 22.58 dB PSNR - Add 9B KV-cache image editing smoke test with verified golden reference image
amepas
force-pushed
the
onboarding-imageedit-kv-flux2klein
branch
from
August 26, 2026 16:45
f04bf0e to
af7a73f
Compare
amepas
marked this pull request as ready for review
August 26, 2026 16:53
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Speeding up Image Editing on Flux2.Klein models by onboarding Flux2.Klein-9B-KV model variant. This is a separate model (same architecture, different weights) that was trained to support caching KV for reference images during image editing.
Builds on PR 461 so some overlapping code changes are included in this PR.
Usage
Just adding the use_kv=True flag to the 9B model call will switch to the KV-cache model support.
Speed-Ups
Block-sizes can be tuned further for each setting. Using default settings still gives substantial improvements
Visual Verification
Prompt: change the painting so she is facing forward instead of looking over her shoulder
Original:
Edited:
Correctness
KV-Cache implementation does not cause any regression or change in behavior on other flux2.klein pipelines.
Below numbers generated on v6-4.
Note: KV-Cache requires a separate model checkpoint/training setup and is only supported on the 9B model.